import torch
import os
from loguru import logger


class Config:
    """
     LeNet5实现图片分类网络模型
    """

    def __init__(self):
        # 数据集相关配置
        self.project_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        self.dataset_dir = os.path.join(self.project_dir, 'data')

        #  模型相关配置
        self.learning_rate = 0.01
        self.batch_size = 128
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        self.epochs = 20
        self.train_info_per_batch = 50
        self.model_save_per_epoch = 2

        # 模型存储路径
        self.model_save_dir = os.path.join(self.project_dir, 'cache')
        if not os.path.exists(self.model_save_dir):
            os.makedirs(self.model_save_dir)
        self.model_save_path = os.path.join(self.model_save_dir, 'mode.pkl')

        logger.info("=" * 20 + " LeNet5 Configuration " + "=" * 20)
        for key, value in self.__dict__.items():
            logger.info(f"{key} = {value}")


cfg = Config()
